import os
import time
import subprocess
import shutil
import logging
from pathlib import Path

# 配置
BASE_MIN = 30 * 60  # 基础预算阶段时间（秒），例如30分钟
EXTENDED = 6 * 60 * 60  # 延长阶段时间（秒），例如6小时
STAGNATION_WINDOW = 6 * 60 * 60  # 覆盖停滞时间（秒），例如6小时
POLL_INTERVAL = 30  # 每30秒检查一次 fuzzer_stats

# 目录配置
SEEDS_DIR = "ai_seeds"  # 种子目录
OUTPUT_DIR = "afl_out"  # AFL++ 输出目录
TARGET_PROGRAMS = ["test_programs/vulnerable_app.c"]  # 目标程序列表，可以按需求调整

# 初始化日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')


def minimize_seed(crash_dir, minimized_seed_dir, target_program):
    """
    使用 afl-cmin 对崩溃样本进行最小化
    """
    afl_cmin_command = [
        "afl-cmin",
        "-i", crash_dir,  # 崩溃样本所在目录
        "-o", minimized_seed_dir,  # 最小化后的种子保存目录
        "--", target_program, "@@"  # 目标程序及输入占位符
    ]
    logging.info(f"Running seed minimization on {crash_dir}...")
    subprocess.run(afl_cmin_command)
    logging.info(f"Minimized seeds saved to {minimized_seed_dir}")

# 函数：读取 fuzzer_stats
def read_fuzzer_stats(output_dir):
    stats_file = os.path.join(output_dir, "fuzzer_stats")
    stats = {}
    if os.path.exists(stats_file):
        with open(stats_file, 'r') as file:
            for line in file:
                if ":" in line:
                    key, value = line.split(":", 1)
                    stats[key.strip()] = value.strip()
    return stats


# 函数：启动 AFL++ 模糊测试
def run_afl_fuzz(target_program, seed_dir, output_dir):
    afl_command = [
        "afl-fuzz",
        "-i", seed_dir, "-o", output_dir,
        "--", target_program, "@@"
    ]
    logging.info(f"Starting AFL++ for {target_program}...")
    return subprocess.Popen(afl_command)


# 函数：停止 AFL++ 进程
def stop_afl_process(process):
    if process.poll() is None:
        process.terminate()
    process.wait()


# 函数：检查 AFL++ 是否符合停止准则
def check_stop_conditions(stats, start_time, corpus_count_last, crash_last):
    """
    检查 AFL++ 的路径覆盖与崩溃数，判断是否满足停止条件。
    :param stats: fuzzer_stats 中的统计数据
    :param start_time: AFL++ 开始运行的时间
    :param corpus_count_last: 上次检查时的种子数
    :param crash_last: 上次检查时的崩溃数
    :return: 是否达到停止条件（True/False），停止原因
    """
    # 获取 corpus_count 和 saved_crashes
    corpus_count = int(stats.get("corpus_count", 0))
    saved_crashes = int(stats.get("saved_crashes", 0))

    # 检查覆盖数是否停滞
    if corpus_count == corpus_count_last:
        logging.info(f"Corpus count stagnated, corpus_count: {corpus_count}")
        if time.time() - start_time >= BASE_MIN:
            return True, "Corpus count stagnation"

    # 检查崩溃数是否增加
    if saved_crashes > crash_last:
        logging.info(f"Crash found: {saved_crashes}")
        return True, "Crash found"

    return False, ""


# 函数：开始模糊测试并处理任务
def fuzz_scheduler():
    for target_program in TARGET_PROGRAMS:
        target_name = Path(target_program).stem
        seed_dir = os.path.join(SEEDS_DIR, target_name)
        output_dir = os.path.join(OUTPUT_DIR, target_name)

        os.makedirs(output_dir, exist_ok=True)
        start_time = time.time()
        corpus_count_last = 0
        crash_last = 0

        # 启动 AFL++ 并开始测试
        afl_process = run_afl_fuzz(target_program, seed_dir, output_dir)

        # 监控 AFL++ 进程
        while True:
            # 每 30 秒读取一次 fuzzer_stats
            stats = read_fuzzer_stats(output_dir)
            stop_condition_met, reason = check_stop_conditions(stats, start_time, corpus_count_last, crash_last)

            if stop_condition_met:
                logging.info(f"Stopping fuzzing due to: {reason}")
                stop_afl_process(afl_process)
                # 当发现崩溃时，执行复现与种子最小化操作
                if reason == "Crash found":
                    crash_dir = os.path.join(output_dir, "crashes")
                    minimized_seed_dir = os.path.join(output_dir, "minimized_seeds")
                    # Step 1: 保存崩溃种子到复现队列
                    logging.info(f"Adding {target_name} to crash queue for reproduction.")
                    # 这里你可以把崩溃样本移到高优先级队列，或者调用其他复现/验证工具
                    # Step 2: 执行种子最小化
                    minimize_seed(crash_dir, minimized_seed_dir, target_program)
                else:
                    # 如果覆盖停滞，延长测试时间或做种子最小化
                    logging.info(f"No progress in {target_name}, moving to low priority.")
                break

            # 如果崩溃数有增加，记录并处理
            if reason == "Crash found":
                crash_last = int(stats.get("saved_crashes", 0))

            # 检查是否进入 extended 阶段
            elapsed_time = time.time() - start_time
            if elapsed_time >= BASE_MIN and corpus_count_last != int(stats.get("corpus_count", 0)):
                logging.info(f"Extending fuzzing for {target_name} to {EXTENDED / 3600} hours.")
                start_time = time.time()  # 重置时间计数
                corpus_count_last = int(stats.get("corpus_count", 0))

            time.sleep(POLL_INTERVAL)  # 每 30 秒轮询一次 fuzzer_stats


# 运行调度器
if __name__ == "__main__":
    fuzz_scheduler()
